import numbers

import mindspore as ms
from mindspore import nn
from mindspore import ops
from PIL import ImageEnhance, Image
import numpy as np


def _is_pil_image(img):
    return isinstance(img, Image.Image)

class ColorJitter:
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        super().__init__()
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)


    def forward(self, img):
        # randperm = ops.Randperm(max_length=4)  # only support GPU
        # n = ms.Tensor([4], dtype=ms.int32)
        # fn_idx = randperm(n)
        fn_idx = np.arange(4)
        shuffle_flag = np.random.shuffle(fn_idx)
        for fn_id in fn_idx:
            if fn_id == 0 and self.brightness is not None:
                brightness = self.brightness
                # brightness_factor = ms.Tensor(1.0).uniform_(brightness[0], brightness[1]).item()
                brightness_factor = np.random.uniform(brightness[0], brightness[1])
                enhancer = ImageEnhance.Brightness(img)
                img = enhancer.enhance(brightness_factor)

            if fn_id == 1 and self.contrast is not None:
                contrast = self.contrast
                # contrast_factor = ms.Tensor(1.0).uniform_(contrast[0], contrast[1]).item()
                contrast_factor = np.random.uniform(contrast[0], contrast[1])
                enhancer = ImageEnhance.Contrast(img)
                img = enhancer.enhance(contrast_factor)

            if fn_id == 2 and self.saturation is not None:
                saturation = self.saturation
                # saturation_factor = ms.Tensor(1.0).uniform_(saturation[0], saturation[1]).item()
                saturation_factor = np.random.uniform(saturation[0], saturation[1])
                enhancer = ImageEnhance.Color(img)
                img = enhancer.enhance(saturation_factor)

            if fn_id == 3 and self.hue is not None:
                hue = self.hue
                # hue_factor = ms.Tensor(1.0).uniform_(hue[0], hue[1]).item()
                hue_factor = np.random.uniform(hue[0], hue[1])

                # img = F.adjust_hue(img, hue_factor)
                if not (-0.5 <= hue_factor <= 0.5):
                    raise ValueError('hue_factor ({}) is not in [-0.5, 0.5].'.format(hue_factor))

                if not _is_pil_image(img):
                    raise TypeError('img should be PIL Image. Got {}'.format(type(img)))

                input_mode = img.mode
                if input_mode in {'L', '1', 'I', 'F'}:
                    return img

                h, s, v = img.convert('HSV').split()

                np_h = np.array(h, dtype=np.uint8)
                # uint8 addition take cares of rotation across boundaries
                with np.errstate(over='ignore'):
                    np_h += np.uint8(hue_factor * 255)
                h = Image.fromarray(np_h, 'L')

                img = Image.merge('HSV', (h, s, v)).convert(input_mode)

        return img

    def _check_input(self, value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
            value = [center - float(value), center + float(value)]
            if clip_first_on_zero:
                value[0] = max(value[0], 0.0)
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError("{} values should be between {}".format(name, bound))
        else:
            raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name))

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value